"""
@FileName：cuda.py
@Author：Huterox
@Description：Go For It
@Time：2024/4/4 16:48
@Copyright：©2018-2024 awesome!
"""

import numpy as np
from mini_torch import Variable
# 默认开启GPU支持
gpu_enable = True

try:
    import cupy as cp
    cupy = cp
except ImportError:
    gpu_enable = False


def get_array_module(x):
    """
    根据输入的数组类型返回相应的数组模块，用于判断应该使用 NumPy 还是 CuPy。

    参数:
        x (Variable 或 numpy.ndarray 或 cupy.ndarray): 用于确定使用 NumPy 还是 CuPy 的值。

    返回值:
        module: 根据参数返回 `cupy` 或 `numpy` 模块。
    """
    if isinstance(x, Variable):
        x = x.data  # 如果x是Variable类型，则获取其数据属性

    if not gpu_enable:  # 如果GPU未启用
        return np  # 返回NumPy模块
    xp = cp.get_array_module(x)  # 获取数组模块
    return xp  # 返回获取到的数组模块


def as_numpy(x):
    """
    将输入对象转换为 `numpy.ndarray`。

    参数:
        x (numpy.ndarray 或 cupy.ndarray): 可以转换为 `numpy.ndarray` 的任意对象。

    返回值:
        `numpy.ndarray`: 转换后的数组。
    """
    if isinstance(x, Variable):
        x = x.data  # 如果x是Variable类型，则获取其数据属性
    if np.isscalar(x):  # 如果x是标量
        return np.array(x)  # 将x转换为NumPy数组
    elif isinstance(x, np.ndarray):  # 如果x已经是NumPy数组
        return x  # 直接返回x
    return cp.asnumpy(x)  # 如果x是CuPy数组，则转换为NumPy数组


def as_cupy(x):
    """
    将输入对象转换为 `cupy.ndarray`。

    参数:
        x (numpy.ndarray 或 cupy.ndarray): 可以转换为 `cupy.ndarray` 的任意对象。

    返回值:
        `cupy.ndarray`: 转换后的数组。
    """
    if isinstance(x, Variable):
        x = x.data  # 如果x是Variable类型，则获取其数据属性

    if not gpu_enable:  # 如果GPU未启用
        raise Exception('Cupy 无法加载')  # 抛出异常
    return cp.asarray(x)  # 将x转换为CuPy数组并返回